import math
import numpy as np
import matplotlib.pyplot as plt

h = 0.02    # 空间步长
N = 750     # 空间步数
dt = 0.0002  # 时间步长
M = 60      # 时间的步数
R = np.zeros([N + 1, M + 1])  # 储存迭代值
I = np.zeros([N + 1, M + 1])
D = np.zeros([N + 1, M + 1])  # 密度
time_ls = np.arange(0, (M + 1) * dt, dt)

# 初始条件
for i in range(0, N):
    R[i, 0] = pow(math.e, -(i * h - 5) ** 2)
    I[i, 0] = 0

# 递推关系
for k in range(0, M):  # 时间
    for i in range(1, N):  # 空间
        R[i, k + 1] = R[i, k] - 0.5 * (I[i + 1, k] - 2 * I[i, k] + I[i - 1, k])
        I[i, k + 1] = I[i, k] + 0.5 * (R[i + 1, k] - 2 * R[i, k] + R[i - 1, k])
        D[i, k + 1] = R[i, k + 1] ** 2 + I[i, k + 1] ** 2

# 不同位置的温度随时间变化
plt.plot(time_ls, D[200, :], 'g-', label='x=4', linewidth=1.0)
plt.plot(time_ls, D[250, :], 'b-', label='x=5', linewidth=1.0)
plt.plot(time_ls, D[300, :], 'k-', label='x=6', linewidth=1.0)
plt.plot(time_ls, D[350, :], 'r-', label='x=7', linewidth=1.0)
plt.plot(time_ls, D[400, :], 'y-', label='x=8', linewidth=1.0)
plt.ylabel('u(x,t)', fontsize=20)
plt.xlabel('t', fontsize=20)
plt.xlim(0, 0.012)
plt.ylim(0, 0.2)
plt.legend(loc='upper left')
plt.savefig('几率密度时间变化.jpg')
plt.show()
